//
// Created by DELL on 2022/4/21.
//

#include "Shader.h"
#include "../EventSystem.h"
#include "../TimeSys.h"
#include "vulkan/vk_enum_string_helper.h"
#include <iomanip>

class ShaderIncluder : public glslang::TShader::Includer
{
public:
    IncludeResult* includeLocal(const char* headerName, const char* includerName, size_t inclusionDepth) override
    {
        auto fileloaded = readFile(headerName);

        if (fileloaded.empty())
        {
            return nullptr;
        }

        auto content = new char[fileloaded.size()];
        std::memcpy(content, fileloaded.data(), fileloaded.size());
        return new IncludeResult(headerName, content, fileloaded.size(), content);
    }

    IncludeResult* includeSystem(const char* headerName, const char* includerName, size_t inclusionDepth) override
    {
        auto fileloaded = readFile(headerName);

        if (fileloaded.empty())
        {
            return nullptr;
        }

        auto content = new char[fileloaded.size()];
        std::memcpy(content, fileloaded.data(), fileloaded.size());
        return new IncludeResult(headerName, content, fileloaded.size(), content);
    }

    void releaseInclude(IncludeResult* result) override
    {
        if (result)
        {
            delete[] static_cast<char*>(result->userData);
            delete result;
        }
    }
};
BEGIN_NS_VULKAN

bool Shader::isInitGlslang {false};
CustomEventListenerRef Shader::systemExitEvent;

ShaderRef Shader::create(const std::vector<ShaderInfo>& shaderInfos)
{
    if (!Shader::isInitGlslang)
    {
        if (!glslang::InitializeProcess())
        {
            throw std::runtime_error("Failed to initialize glslang process");
        }
        Shader::isInitGlslang = true;
        Shader::systemExitEvent = CustomEventListener::create(CustomEvent::exitSystemEvent);
        Shader::systemExitEvent->onCustomEvent = [](auto e){
            glslang::FinalizeProcess();
        };
        EventSystem::get()->subscribe(Shader::systemExitEvent);
    }
    auto shaderRef = new Shader;

    std::vector<std::shared_ptr<glslang::TShader>> shaders(shaderInfos.size());
    for(int i = 0; i < shaderInfos.size(); ++i)
    {
//        TimeSys::startTiming();
        if (shaderInfos[i].isPath)
        {
            shaders[i] = createShader(readFileToString(shaderInfos[i].source), stage2ShLanguage(shaderInfos[i].stage));
        }
        else
        {
            shaders[i] = createShader(shaderInfos[i].source, stage2ShLanguage(shaderInfos[i].stage));
        }
//        LOGI("Read Shader time: {}", TimeSys::timing());
    }
    auto program = createProgram(shaders);
    shaderRef->initDescriptorSetLayout(program.get());

    glslang::SpvOptions spvOptions;
    spvOptions.generateDebugInfo = false;
    spvOptions.disableOptimizer = false;
    spvOptions.optimizeSize = true;
    spv::SpvBuildLogger logger;

    for(const auto & shaderInfo : shaderInfos)
    {
        shaderRef->stages.push_back(shaderInfo.stage);
        std::vector<uint32_t> spirv;
        GlslangToSpv(*program->getIntermediate(stage2ShLanguage(shaderInfo.stage)), spirv, &logger, &spvOptions);
        shaderRef->spirvCodes.push_back(std::move(spirv));
    }
    return ShaderRef(shaderRef);
}

ShaderRef Shader::create(const std::string& vertSrc, const std::string& fragSrc)
{
    return create({ ShaderInfo{VK_SHADER_STAGE_VERTEX_BIT, vertSrc, false},
                   ShaderInfo{VK_SHADER_STAGE_FRAGMENT_BIT, fragSrc, false}});
}

ShaderRef Shader::createByPath(const std::string& vertPath, const std::string& fragPath)
{
    return create({ ShaderInfo{VK_SHADER_STAGE_VERTEX_BIT, vertPath},
                   ShaderInfo{VK_SHADER_STAGE_FRAGMENT_BIT, fragPath}});
}

std::shared_ptr<glslang::TShader> Shader::createShader(const std::string& src, EShLanguage language)
{
    auto* shader = new glslang::TShader(language);
    TBuiltInResource resources {};
    populateBuiltInResource(resources);

    auto shaderSourcePtr = src.c_str();
    shader->setStrings(&shaderSourcePtr, 1);

    auto defaultVersion = 430;
    shader->setEnvInput(glslang::EShSourceGlsl, language, glslang::EShClientVulkan, defaultVersion);
    shader->setEnvClient(glslang::EShClientVulkan, glslang::EShTargetVulkan_1_2);
    shader->setEnvTarget(glslang::EShTargetSpv, glslang::EShTargetSpv_1_5);

    ShaderIncluder includer;
    std::string    str;
    // Enable SPIR-V and Vulkan rules when parsing GLSL.
    auto messages = static_cast<EShMessages>(EShMsgSpvRules | EShMsgVulkanRules | EShMsgDefault);

    if (!shader->preprocess(&resources, defaultVersion, ENoProfile, false, false, messages, &str, includer))
    {
        LOGE("{} Shader preprocess error !!\n{} \n{}",language, shader->getInfoLog(), shader->getInfoDebugLog());
    }

    if (!shader->parse(&resources, defaultVersion, true, messages, includer))
    {
        LOGE("{} Shader parse error !!\n{} \n{}",language, shader->getInfoLog(), shader->getInfoDebugLog());
    }
    return std::shared_ptr<glslang::TShader>(shader);
}

std::shared_ptr<glslang::TProgram> Shader::createProgram(const std::vector<std::shared_ptr<glslang::TShader>>& shaders)
{
    auto messages = static_cast<EShMessages>(EShMsgSpvRules | EShMsgVulkanRules | EShMsgDefault);

    auto* program = new glslang::TProgram();
    for (auto& shader: shaders)
    {
        program->addShader(shader.get());
    }

    if (!program->link(messages) || !program->mapIO())
    {
        LOGE("Shader link error !! \n{} \n{}", program->getInfoLog(), program->getInfoDebugLog());
    }
    program->buildReflection();
    return  std::shared_ptr<glslang::TProgram>(program);
}

void Shader::initDescriptorSetLayout(glslang::TProgram* program)
{
//    for (int32_t i = 0; i < program->getNumLiveAttributes(); i++)
//    {
//        const auto& reflection = program->getPipeInput(i);
//        auto& qualifier = reflection.getType()->getQualifier();
//        auto index = (int32_t)qualifier.layoutLocation;
//        auto name = reflection.name;
//        auto type = glslangType2UniformType(reflection.getType());
//    }

    this->descriptorSetLayout.blocks.reserve(program->getNumLiveUniformBlocks());
    for (int32_t i = program->getNumLiveUniformBlocks() - 1; i >= 0; i--)
    {
        const auto& reflection = program->getUniformBlock(i);
        UniformBlock block;
        block.name = reflection.name;
        block.binding = reflection.getBinding();
        block.size = reflection.size;
        block.stages = glslangStages2Vulkan(reflection.stages);
        auto type = reflection.getType();
//        if (type != nullptr && type->getQualifier().hasPacking())
//        {
//            block.isStd140 = type->getQualifier().layoutPacking == glslang::ElpStd140;
//        }
        this->descriptorSetLayout.blocks.push_back(block);
    }

    for (int32_t i = 0; i < program->getNumLiveUniformVariables(); i++)
    {
        const auto& reflection = program->getUniform(i);
        auto pos = reflection.name.find('.');
        auto type = reflection.getType();
        if (pos != std::string::npos)
        {
            auto blockName = reflection.name.substr(0, pos);
            for (auto& block : this->descriptorSetLayout.blocks)
            {
                if (blockName == block.name)
                {
                    block.fields[type->getFieldName().c_str()] = UniformField{
                        glslangType2UniformType(type), reflection.offset
                    };
                    break;
                }
            }
        }
        else if (type->getBasicType() == glslang::EbtSampler)
        {
            descriptorSetLayout.samples[reflection.name] = { reflection.name, kU_sampler,
                                                            reflection.getBinding(), glslangStages2Vulkan(reflection.stages)};
        }
        else if (type->isAtomic())
        {
            descriptorSetLayout.samples[reflection.name] = { reflection.name, kU_atomic,
                                                            reflection.getBinding() , glslangStages2Vulkan(reflection.stages)};
        }
    }
}

void Shader::populateBuiltInResource(TBuiltInResource& resources)
{
    resources.maxLights = 32;
    resources.maxClipPlanes = 6;
    resources.maxTextureUnits = 32;
    resources.maxTextureCoords = 32;
    resources.maxVertexAttribs = 64;
    resources.maxVertexUniformComponents = 4096;
    resources.maxVaryingFloats = 64;
    resources.maxVertexTextureImageUnits = 32;
    resources.maxCombinedTextureImageUnits = 80;
    resources.maxTextureImageUnits = 32;
    resources.maxFragmentUniformComponents = 4096;
    resources.maxDrawBuffers = 32;
    resources.maxVertexUniformVectors = 128;
    resources.maxVaryingVectors = 8;
    resources.maxFragmentUniformVectors = 16;
    resources.maxVertexOutputVectors = 16;
    resources.maxFragmentInputVectors = 15;
    resources.minProgramTexelOffset = -8;
    resources.maxProgramTexelOffset = 7;
    resources.maxClipDistances = 8;
    resources.maxComputeWorkGroupCountX = 65535;
    resources.maxComputeWorkGroupCountY = 65535;
    resources.maxComputeWorkGroupCountZ = 65535;
    resources.maxComputeWorkGroupSizeX = 1024;
    resources.maxComputeWorkGroupSizeY = 1024;
    resources.maxComputeWorkGroupSizeZ = 64;
    resources.maxComputeUniformComponents = 1024;
    resources.maxComputeTextureImageUnits = 16;
    resources.maxComputeImageUniforms = 8;
    resources.maxComputeAtomicCounters = 8;
    resources.maxComputeAtomicCounterBuffers = 1;
    resources.maxVaryingComponents = 60;
    resources.maxVertexOutputComponents = 64;
    resources.maxGeometryInputComponents = 64;
    resources.maxGeometryOutputComponents = 128;
    resources.maxFragmentInputComponents = 128;
    resources.maxImageUnits = 8;
    resources.maxCombinedImageUnitsAndFragmentOutputs = 8;
    resources.maxCombinedShaderOutputResources = 8;
    resources.maxImageSamples = 0;
    resources.maxVertexImageUniforms = 0;
    resources.maxTessControlImageUniforms = 0;
    resources.maxTessEvaluationImageUniforms = 0;
    resources.maxGeometryImageUniforms = 0;
    resources.maxFragmentImageUniforms = 8;
    resources.maxCombinedImageUniforms = 8;
    resources.maxGeometryTextureImageUnits = 16;
    resources.maxGeometryOutputVertices = 256;
    resources.maxGeometryTotalOutputComponents = 1024;
    resources.maxGeometryUniformComponents = 1024;
    resources.maxGeometryVaryingComponents = 64;
    resources.maxTessControlInputComponents = 128;
    resources.maxTessControlOutputComponents = 128;
    resources.maxTessControlTextureImageUnits = 16;
    resources.maxTessControlUniformComponents = 1024;
    resources.maxTessControlTotalOutputComponents = 4096;
    resources.maxTessEvaluationInputComponents = 128;
    resources.maxTessEvaluationOutputComponents = 128;
    resources.maxTessEvaluationTextureImageUnits = 16;
    resources.maxTessEvaluationUniformComponents = 1024;
    resources.maxTessPatchComponents = 120;
    resources.maxPatchVertices = 32;
    resources.maxTessGenLevel = 64;
    resources.maxViewports = 16;
    resources.maxVertexAtomicCounters = 0;
    resources.maxTessControlAtomicCounters = 0;
    resources.maxTessEvaluationAtomicCounters = 0;
    resources.maxGeometryAtomicCounters = 0;
    resources.maxFragmentAtomicCounters = 8;
    resources.maxCombinedAtomicCounters = 8;
    resources.maxAtomicCounterBindings = 1;
    resources.maxVertexAtomicCounterBuffers = 0;
    resources.maxTessControlAtomicCounterBuffers = 0;
    resources.maxTessEvaluationAtomicCounterBuffers = 0;
    resources.maxGeometryAtomicCounterBuffers = 0;
    resources.maxFragmentAtomicCounterBuffers = 1;
    resources.maxCombinedAtomicCounterBuffers = 1;
    resources.maxAtomicCounterBufferSize = 16384;
    resources.maxTransformFeedbackBuffers = 4;
    resources.maxTransformFeedbackInterleavedComponents = 64;
    resources.maxCullDistances = 8;
    resources.maxCombinedClipAndCullDistances = 8;
    resources.maxSamples = 4;
    resources.limits.nonInductiveForLoops = true;
    resources.limits.whileLoops = true;
    resources.limits.doWhileLoops = true;
    resources.limits.generalUniformIndexing = true;
    resources.limits.generalAttributeMatrixVectorIndexing = true;
    resources.limits.generalVaryingIndexing = true;
    resources.limits.generalSamplerIndexing = true;
    resources.limits.generalVariableIndexing = true;
    resources.limits.generalConstantMatrixVectorIndexing = true;
}

EShLanguage Shader::stage2ShLanguage(VkShaderStageFlagBits stage)
{
    switch (stage)
    {
        case VK_SHADER_STAGE_VERTEX_BIT: return EShLangVertex;
        case VK_SHADER_STAGE_TESSELLATION_CONTROL_BIT: return EShLangTessControl;
        case VK_SHADER_STAGE_TESSELLATION_EVALUATION_BIT: return EShLangTessEvaluation;
        case VK_SHADER_STAGE_GEOMETRY_BIT: return EShLangGeometry;
        case VK_SHADER_STAGE_FRAGMENT_BIT: return EShLangFragment;
        case VK_SHADER_STAGE_COMPUTE_BIT: return EShLangCompute;
        default:
            break;
    }
    return EShLangCount;
}

UniformType Shader::glslangType2UniformType(const glslang::TType* type)
{
    switch (type->getBasicType())
    {
        case glslang::EbtBool: return kU_bool;
        case glslang::EbtInt8: return kU_int8;
        case glslang::EbtInt16: return kU_int16;
        case glslang::EbtInt: return kU_int32;
        case glslang::EbtUint8: return kU_uint8;
        case glslang::EbtUint16: return kU_uint16;
        case glslang::EbtUint: return kU_uint32;
        case glslang::EbtDouble: return kU_double;
        case glslang::EbtAtomicUint: return kU_atomic;
        case glslang::EbtStruct: return kU_struct;
        case glslang::EbtFloat:
        {
            switch (type->getVectorSize())
            {
                case 0:
                {
                    if (type->getMatrixCols() == 3 && type->getMatrixRows() == 3)
                    {
                        return kU_mat3;
                    }
                    else if (type->getMatrixCols() == 4 && type->getMatrixRows() == 4)
                    {
                        return kU_mat4;
                    }
                    return kU_void;
                }
                case 1: return kU_float;
                case 2: return KU_vec2;
                case 3: return KU_vec3;
                case 4: return KU_vec4;
                default: break;
            }
        }
        case glslang::EbtSampler: return kU_sampler;
        default: break;
    }

    return kU_void;
}

VkShaderStageFlags Shader::glslangStages2Vulkan(EShLanguageMask stages)
{
    constexpr EShLanguageMask masks[] {EShLangVertexMask, EShLangTessControlMask, EShLangTessEvaluationMask,
        EShLangGeometryMask, EShLangFragmentMask, EShLangComputeMask, EShLangRayGenMask};

    VkShaderStageFlags flags{0};
    for (auto mask : masks)
    {
        if ((mask & stages) > 0)
        {
            flags |= glslangStageBit2Vulkan(mask);
        }
    }
    return flags;
}

VkShaderStageFlagBits Shader::glslangStageBit2Vulkan(EShLanguageMask stage)
{
    switch (stage)
    {
        case EShLangVertexMask: return VK_SHADER_STAGE_VERTEX_BIT;
        case EShLangTessControlMask: return VK_SHADER_STAGE_TESSELLATION_CONTROL_BIT;
        case EShLangTessEvaluationMask: return VK_SHADER_STAGE_TESSELLATION_EVALUATION_BIT;
        case EShLangGeometryMask: return VK_SHADER_STAGE_GEOMETRY_BIT;
        case EShLangFragmentMask: return VK_SHADER_STAGE_FRAGMENT_BIT;
        case EShLangComputeMask: return VK_SHADER_STAGE_COMPUTE_BIT;
        case EShLangRayGenMask: return VK_SHADER_STAGE_RAYGEN_BIT_KHR;
        default:
            break ;
    }
    return (VkShaderStageFlagBits)0;
}
void Shader::dumpSpirv(const std::string& fileName)
{
    std::ofstream file(fileName, std::ios::binary|std::ios::out|std::ios::trunc);

    file.setf(std::ios::hex);
    for (int i = 0; i < this->stages.size(); ++i)
    {
        file << "// " << string_VkShaderStageFlagBits(stages[i]) << "\n";
        for (int j = 0; j < this->spirvCodes[i].size(); ++j)
        {

            file << "0x" << std::setw(8) << std::setfill('0') << std::setbase(16) << this->spirvCodes[i][j] << ", ";
            if (j % 10 == 9)
            {
                file << "\n";
            }
        }
        file << "\n\n";
    }

}

END_NS_VULKAN
